package edu.utdallas.cs.securechat.main;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.util.Date;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

public class ServerProtocol {
	SecretKey secretKey;
	Key publicKey;
	Key privateKey;
	KeyStore ks;
	public static long expiresAfter = 360000000;
	final int keysize = 16;
	
	public ServerProtocol() {
		super();
	}
	
	public Key getPublicKey()
	{
		return publicKey;
	}
	
	public ServerProtocol(SecretKey secret, Key pubKey, Key privKey)
	{
		this.secretKey = secret;
		this.publicKey = pubKey;
		this.privateKey = privKey;
	}
	
	public void setSecretKey(SecretKey key)
	{
		this.secretKey = key;
	}
	
	public void setSecretKey(byte[] key){
		this.secretKey = new SecretKeySpec(key, "AES");
	}
	
	public String hashPassword(byte[] pass) throws NoSuchAlgorithmException{
		MessageDigest hash = MessageDigest.getInstance("SHA-1");
		hash.reset();
	    hash.update(pass);
	    //hash.digest();
        String s = new BigInteger(1, hash.digest()).toString(16);
        return (s.length() % 2 == 0) ? s : "0" + s;
	}
	
	public byte[] RSADecrypt(byte[] message) throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException
	{
	    Cipher cipher = Cipher.getInstance("RSA");
	    cipher.init(Cipher.DECRYPT_MODE, privateKey);
	    byte[] plainText = cipher.doFinal(message);
	    return plainText;
	}
		
	public byte[] AESEncrypt(byte[] message, SecretKey key) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException
	{
		Cipher c = Cipher.getInstance("AES");
		SecretKeySpec k = new SecretKeySpec(key.getEncoded(), "AES");
		c.init(Cipher.ENCRYPT_MODE, k);
		byte[] encryptedData = c.doFinal(message);
		return encryptedData;
	}
		
	public byte[] AESDecrypt(byte[] message, SecretKey key) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException
	{
		Cipher c = Cipher.getInstance("AES");
		SecretKeySpec k = new SecretKeySpec(key.getEncoded(), "AES");
		c.init(Cipher.DECRYPT_MODE, k);
		byte[] data = c.doFinal(message);
		return data;	
	}
	
	public byte[] grantTGT(SecretKey clientKey, String username) throws IOException, NoSuchAlgorithmException, InvalidKeyException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException
	{
		long time = new Date().getTime();
		long e = time + ServerProtocol.expiresAfter;

		ByteArrayOutputStream response = new ByteArrayOutputStream();  
		DataOutputStream dos = new DataOutputStream(response);  
		dos.writeLong(e);  
		
		response.write(this.generateAESKey().getEncoded());
		byte[] r = response.toByteArray();

		ByteArrayOutputStream ticketStream = new ByteArrayOutputStream();
		ticketStream.write(r);
		ticketStream.write(username.getBytes());
		byte[] t = ticketStream.toByteArray();
		ByteArrayOutputStream result = new ByteArrayOutputStream();
		result.write(this.AESEncrypt(r, clientKey));
		result.write(this.AESEncrypt(t, this.secretKey));
		
		return result.toByteArray();
	}
	
	public TicketRequest handleTicketRequest(byte[] ticketRequest) throws Exception{
		ByteArrayInputStream ticketRequestStream = new ByteArrayInputStream(ticketRequest);
		int l = (int)ticketRequestStream.read();
		byte[] Erequest = new byte[l];
		byte[] TGT = new byte[ticketRequestStream.available()-Erequest.length];
		ticketRequestStream.read(Erequest);
		ticketRequestStream.read(TGT);
		ByteArrayInputStream TGTstream = new ByteArrayInputStream(this.AESDecrypt(TGT, this.secretKey));
		
		
		DataInputStream e = new DataInputStream(TGTstream);
		long expire = e.readLong();
		byte[] sKey = new byte[keysize];
		TGTstream.read(sKey);
		byte[] u = new byte[TGTstream.available()];
		TGTstream.read(u);
		
		String user1 = new String(u);
		
		if (expire >= new Date().getTime()){
			SecretKey sessionKey = new SecretKeySpec(sKey, "AES");
			ByteArrayInputStream requestStream = new ByteArrayInputStream(this.AESDecrypt(Erequest, sessionKey));
			u = new byte[requestStream.available()-8];
			e = new DataInputStream(requestStream);
			long timestamp = e.readLong();
			requestStream.read(u);
			String user2 = new String(u);
			
			return new TicketRequest(user1, user2, sessionKey, timestamp);
		} else {
			throw new Exception("Ticket is expired");
		}
	}
	
	public byte[] grantTicket(TicketRequest request, SecretKey user2Key) throws IOException, NoSuchAlgorithmException, InvalidKeyException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException{
		ByteArrayOutputStream requestStream = new ByteArrayOutputStream();  
		DataOutputStream dos = new DataOutputStream(requestStream);
		dos.writeLong(request.getTimestamp()+1);  
		long expire = new Date().getTime() + ServerProtocol.expiresAfter;
		dos.writeLong(expire);  

		byte[] key = this.generateAESKey().getEncoded();
		requestStream.write(key);
		requestStream.write(request.getUser2().getBytes());
				
		ByteArrayOutputStream authStream = new ByteArrayOutputStream();
		dos = new DataOutputStream(authStream);
		dos.writeLong(expire);
		authStream.write(key);

		authStream.write(request.getUser1().getBytes());
		ByteArrayOutputStream resultStream = new ByteArrayOutputStream();
		byte[] Eauth = this.AESEncrypt(authStream.toByteArray(), user2Key);
		resultStream.write(Eauth.length);
		resultStream.write(Eauth);
		resultStream.write(this.AESEncrypt(requestStream.toByteArray(), request.getKey()));
		
		return resultStream.toByteArray();
		
	}
	
	public void generateRSAKeys() throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException
	{
	    SecureRandom random = new SecureRandom();
	    KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");

	    generator.initialize(512, random);

	    KeyPair pair = generator.generateKeyPair();
	    this.publicKey = pair.getPublic();
	    this.privateKey = pair.getPrivate();
	}
	
	
	
	public SecretKey generateAESKey() throws NoSuchAlgorithmException
	{
		KeyGenerator generator = KeyGenerator.getInstance("AES");
		generator.init(128);
		SecretKey skey = generator.generateKey();
		return skey;
	}
	
	public SecretKey keyFromPassword(char[] password) throws NoSuchAlgorithmException, InvalidKeySpecException
	{
		SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
		KeySpec spec = new PBEKeySpec(password, new byte[]{4, 8, 15, 16, 23, 42}, 65536, 128);
		SecretKey tmp = factory.generateSecret(spec);
		SecretKey secret = new SecretKeySpec(tmp.getEncoded(), "AES");
		return secret;
	}
}
